import winreg
import os
import sys
import ctypes
from tkinter import messagebox

def is_admin():
    """检查是否具有管理员权限"""
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False

def register_file_association():
    """注册文件关联"""
    try:
        # 获取当前程序的完整路径
        if getattr(sys, 'frozen', False):
            exe_path = sys.executable  # 打包后的exe路径
        else:
            exe_path = sys.executable
            script_path = os.path.abspath("notepad.py")
        
        # 注册.txt文件类型的打开方式
        with winreg.CreateKey(winreg.HKEY_CLASSES_ROOT, ".txt") as key:
            # 添加到现有的打开方式列表中
            with winreg.CreateKey(key, "OpenWithProgids") as subkey:
                winreg.SetValueEx(subkey, "MJXNotepad.txt", 0, winreg.REG_NONE, b"")

        # 注册应用程序信息
        with winreg.CreateKey(winreg.HKEY_CLASSES_ROOT, "MJXNotepad.txt") as key:
            winreg.SetValue(key, "", winreg.REG_SZ, "码界轩记事本文档")
            
            # 设置图标
            with winreg.CreateKey(key, "DefaultIcon") as subkey:
                if getattr(sys, 'frozen', False):
                    icon_path = exe_path  # 使用exe自身作为图标
                else:
                    icon_path = os.path.abspath("mjx128.ico")
                winreg.SetValue(subkey, "", winreg.REG_SZ, icon_path)
            
            # 设置打开命令
            with winreg.CreateKey(key, "shell\\open\\command") as subkey:
                if getattr(sys, 'frozen', False):
                    command = f'"{exe_path}" "%1"'  # 打包后直接使用exe
                else:
                    command = f'"{exe_path}" "{script_path}" "%1"'
                winreg.SetValue(subkey, "", winreg.REG_SZ, command)

        messagebox.showinfo("成功", '已成功添加"码界轩记事本"到txt文件的打开方式列表中！')
        return True
        
    except Exception as e:
        messagebox.showerror("错误", f"注册失败: {str(e)}")
        return False

if __name__ == "__main__":
    # 检查是否具有管理员权限
    if not is_admin():
        # 请求管理员权限重新运行
        ctypes.windll.shell32.ShellExecuteW(None, "runas", sys.executable, " ".join(sys.argv), None, 1)
    else:
        register_file_association() 